#include <rtthread.h>
#include <rtdevice.h>
#include <stdlib.h>
#include <string.h>
#include <stdio.h>

#ifdef RT_USING_MTD
static rt_mtd_t *_mtd = 0;

static int mget(int argc, char **argv)
{
    if (argc == 1)
        return 1;

    _mtd = rt_mtd_get(argv[1]);
    if (_mtd == 0)
        printf("not found\n");

    return 0;
}
MSH_CMD_EXPORT(mget, MTD get device by name);

static int merase(int argc, char **argv)
{
    int b;
    int ret;

    if (_mtd == 0)
        return 1;

    if (argc == 1)
    {
        int blks;

        printf("erase all of the MTD\n");

        blks = _mtd->size / _mtd->block_size;
        for (b = 0; b < blks; b++)
        {
            ret = rt_mtd_block_erase(_mtd, b);
            if (ret != 0)
                break;
        }
    }

    if (argc == 2)
    {
        b = atoi(argv[1]);
        ret = rt_mtd_block_erase(_mtd, b);
    }

    if (ret)
        printf("erase fail @b %d\n", b);

    return ret;
}
MSH_CMD_EXPORT(merase, MTD erase block);

static int mread(int argc, char **argv)
{
    uint8_t *buf;
    int size;
    int addr = 0;
    int ret;

    if (_mtd == 0)
        return 1;

    size = _mtd->sector_size;
    if (argc == 3)
    {
        addr = atoi(argv[1]);
        size = atoi(argv[2]);
    }
    else if (argc == 2)
    {
        addr = atoi(argv[1]);
    }

    buf = rt_malloc(size);
    if (buf)
    {
        int i;

        rt_memset(buf, 0, size);

        printf("read from %d size %d\n", addr, size);
        ret = rt_mtd_read(_mtd, addr, buf, size);
        printf("ret = %d\n", ret);

        for (i = 0; i < ret; i ++)
        {
            if (i%16 == 0)
                printf("\n");

            printf("%02X ", buf[i]);
        }
        printf("\n");
    }
    rt_free(buf);

    return ret;
}
MSH_CMD_EXPORT(mread, MTD read: [FROM] [SIZE]);

static int mwrite(int argc, char **argv)
{
    char ch = '1';
    uint8_t *buf;
    int size;
    int addr = 0;
    int ret;
    int random = 0;

    if (_mtd == 0)
        return 1;

    size = _mtd->sector_size;
    while (argc)
    {
        switch (argc)
        {
        case 2:
            addr = atoi(argv[1]);
            break;
        case 3:
            size = atoi(argv[2]);
            break;
        case 4:
            ch = argv[3][0];
            if (strcmp(argv[3], "-r") == 0)
                random = 1;
            break;
        }
        argc--;
    }

    buf = rt_malloc(size);
    if (buf)
    {
        printf("write to %d size %d\n", addr, size);
        if (random)
        {
            srand(rt_tick_get());
            for (random = 0; random < size; random++)
                buf[random] = rand();
        }
        else
        {
            rt_memset(buf, ch, size);
        }

        ret = rt_mtd_write(_mtd, addr, buf, size);
        printf("ret = %d\n", ret);
    }
    rt_free(buf);

    return ret;
}
MSH_CMD_EXPORT(mwrite, MTD write: [TO] [SIZE] [CHAR/-r]);

static int mroob(int argc, char **argv)
{
    int ret;
    struct mtd_io_desc desc = {0};
    uint8_t oob[80];
    int addr = 0;

    if (_mtd == 0)
        return 1;

    while (argc)
    {
        switch (argc)
        {
        case 2:
            addr = atoi(argv[1]);
            break;
        case 3:
            switch (argv[2][0])
            {
            case 'p':
                desc.mode = MTD_OPM_PLACE_OOB;
                break;
            case 'a':
                desc.mode = MTD_OPM_AUTO_OOB;
                break;
            case 'r':
                desc.mode = MTD_OPM_RAW;
                break;
            }
            break;
        }
        argc--;
    }

    desc.oobbuf = oob;
    desc.ooblen = sizeof(oob);
    ret = rt_mtd_read_oob(_mtd, addr, &desc);
    printf("read oob from %0X, ret %d, ooblen %d\n", addr, ret, desc.oobretlen);

    return 0;
}
MSH_CMD_EXPORT(mroob, MTD read oob: [FROM] [p/a/r]);

static int meccerr(int argc, char **argv)
{
    char errbits = 1;
    uint8_t *buf;
    int size;
    int addr = 0;
    int ret;
    int random;
    struct mtd_io_desc desc = {0};
    uint8_t mask[2] = {0xFE, 0xFC};
    int pos = 0;

    if (_mtd == 0)
        return 1;

    size = _mtd->sector_size;
    while (argc)
    {
        switch (argc)
        {
        case 2:
            addr = atoi(argv[1]);
            break;
        case 3:
            pos = atoi(argv[2]);
            break;
        case 4:
            errbits = atoi(argv[3]);
            break;
        }
        argc--;
    }

    if (errbits < 1 || errbits > 2)
        errbits = 1;
    if (pos < 0 || pos >= size)
        pos = 0;

    buf = rt_malloc(size);
    if (buf)
    {
        printf("write to %d pos %d errbits %d\n", addr, pos, errbits);

        srand(rt_tick_get());
        for (random = 0; random < size; random++)
            buf[random] = rand();

        ret = rt_mtd_write(_mtd, addr, buf, size);
        printf("ret = %d\n", ret);

        desc.datbuf = buf;
        desc.datlen = size;
        desc.mode = MTD_OPM_RAW;
        buf[0] = buf[pos] & mask[errbits - 1];
        ret = rt_mtd_write_oob(_mtd, addr, &desc);
        printf("ret = %d\n", ret);
    }
    rt_free(buf);

    return ret;
}
MSH_CMD_EXPORT(meccerr, MTD emulate ecc err: [ADDR] [POS] [errbits:1/2]);

static int mpcheck(int argc, char **argv)
{
    char errbits = 1;
    uint8_t *buf;
    int size;
    int addr = 0;
    int ret;
    int random;
    int page = 0;
    int num = 1;

    if (_mtd == 0)
        return 1;

    size = _mtd->sector_size;
    while (argc)
    {
        switch (argc)
        {
        case 2:
            page = atoi(argv[1]);
            break;
        case 3:
            num = atoi(argv[2]);
            break;
        }
        argc--;
    }

    buf = rt_malloc(size);
    if (buf)
    {
        printf("write to %d num %d\n", page, num);
        num += page;
        addr = page * _mtd->sector_size;
        for (; page < num; page++)
        {
            srand(rt_tick_get());
            for (random = 0; random < size; random++)
                buf[random] = rand();
            ret = rt_mtd_write(_mtd, addr, buf, size);
            ret = rt_mtd_read(_mtd, addr, buf, size);
            addr += size;
            if (ret != size)
                printf("fail at page %d\n", page);
        }
    }
    rt_free(buf);

    return ret;
}
MSH_CMD_EXPORT(mpcheck, MTD page write and read check: [PAGE] [NUM]);
#endif
